import os
import numpy as np
from .file_utils import count_files
from .process_bar import ProcessBar


class FileProcessor:
    def __init__(self, input_root: str, processing_function, max_amount: int = 20000,
                 output: bool = True, output_root: str = '', output_fmt: str = 'npy'):
        """
        初始化 FileProcessor 类。

        参数:
        input_root: str
            数据集的根目录路径。
        output_root: str
            输出目录的根目录路径。
        """
        self.input_root = input_root
        self.processing_function = processing_function
        self.max_amount = max_amount
        self.output_root = output_root
        self.output = output
        self.output_fmt = output_fmt

    def process_files(self):
        """
        处理input_root目录下的文件。

        参数:
        processing_function: function
            一个接受文件路径和身份id作为参数的处理函数。
        """
        for id_count, identity_id in enumerate(os.listdir(self.input_root)):
            if id_count >= self.max_amount:
                break
            identity_path = os.path.join(self.input_root, identity_id)
            files_amount = count_files(identity_path)
            print(f"Processing {identity_id}, {files_amount} files.")
            current_index = 0
            pb = ProcessBar(files_amount)

            if os.path.isdir(identity_path):
                for subfolder in os.listdir(identity_path):
                    subfolder_path = os.path.join(identity_path, subfolder)
                    if os.path.isdir(subfolder_path):
                        index = 0
                        for filename in os.listdir(subfolder_path):
                            file_path = os.path.join(subfolder_path, filename)
                            if self.output:
                                processed_data = self.processing_function(file_path, identity_id)

                                if self.output_fmt == 'npy-list':
                                    for data in processed_data:
                                        output_path = self.get_output_path(file_path, index)
                                        self.write_output(data, output_path)
                                        index = index + 1
                                else:
                                    output_path = self.get_output_path(file_path)
                                    self.write_output(processed_data, output_path)
                            else:
                                self.processing_function(file_path, identity_id)
                            current_index = current_index + 1
                            pb.print(current_index)
            print()

    def get_output_path(self, file_path, index: int = -1):
        """
        根据输入文件路径生成输出文件路径。

        参数:
        file_path: str
            输入文件的路径。

        返回:
        output_path: str
            输出文件的路径。
        """
        if index == -1:
            relative_path_basename = os.path.splitext(os.path.relpath(file_path, self.input_root))[0]
        else:
            fmt_index = f'{index:0>{3}}'
            relative_path_basename = os.path.join(os.path.relpath(os.path.dirname(file_path), self.input_root), fmt_index)

        extension = '.npy' if self.output_fmt == 'npy' or self.output_fmt == 'npy-list' else ''
        output_path = f'{os.path.join(self.output_root, relative_path_basename)}{extension}'

        return output_path

    def write_output(self, data, output_path):
        """
        将处理后的数据写入输出文件。

        参数:
        data: any
            要写入的数据。
        output_path: str
            输出文件的路径。
        """
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        if self.output_fmt == 'npy' or self.output_fmt == 'npy-list':
            np.save(output_path, data)
        else:
            with open(output_path, 'w') as f:
                f.write(data)
